import cv2
import numpy as np
 
 # yolov5中，直接使用opencv的dnn模块进行yolov5的模型文件的推理
classes=["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
        "hair drier", "toothbrush"]
 
 
def infer(img,pred,shape):
    w_ratio=img.shape[1]/shape[0]
    h_ratio=img.shape[0]/shape[1]
    # print(w_ratio,h_ratio)
    confidences=[]
    boxes=[]
    class_ids=[]
    boxes_num=pred.shape[1]
    data=pred[0]
    # print("data_shape:",data.shape)
    for i in range(boxes_num):
        da=data[i]#[box,conf,cls]
        confidence=da[4]
        if confidence>0.6:
            score=da[5:]*confidence
            _,_,_,max_score_index=cv2.minMaxLoc(score)#
            max_cls_id=max_score_index[1]
            if score[max_cls_id]>0.25:
                confidences.append(confidence)
                class_ids.append(max_cls_id)
                x,y,w,h=da[0].item(),da[1].item(),da[2].item(),da[3].item()
                nx=int((x-w/2.0)*w_ratio)
                ny=int((y-h/2.0)*h_ratio)
                nw=int(w*w_ratio)
                nh=int(h*h_ratio)
                boxes.append(np.array([nx,ny,nw,nh]))
 
    indexes=cv2.dnn.NMSBoxes(boxes,confidences,0.25,0.45)
    res_ids=[]
    res_confs=[]
    res_boxes=[]
    for i in indexes:
        res_ids.append(class_ids[i])
        res_confs.append(confidences[i])
        res_boxes.append(boxes[i])
 
    # print(res_ids)
    # print(res_confs)
    # print(res_boxes)
    return res_ids,res_confs,res_boxes
 
def draw_rect(img,ids,confs,boxes):
    for i in range(len(ids)):
        cv2.rectangle(img, boxes[i], (0,0,255), 2)
        cv2.rectangle(img, (boxes[i][0],boxes[i][1]-20),(boxes[i][0]+boxes[i][2],boxes[i][1]), (200, 200, 200), -1)
        cv2.putText(img, classes[ids[i]], (boxes[i][0], boxes[i][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 0, 0))
        cv2.putText(img, str(confs[i]), (boxes[i][0]+60, boxes[i][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 0, 0))
 
    cv2.imwrite("res.jpg",img)
    cv2.imshow('img',img)
    cv2.waitKey()
 
 
if __name__=="__main__":
    import time
    st=time.time()
    shape=(640,640)
    src=cv2.imread(r'E:\pythonProject\yolov5master\data\images\bus.jpg')
    # src=cv2.imread(r'd:\yuanshen.jpg')
    img=src.copy()
    # pip install -U onnx-simplifier --user
    # step1: python export.py --weights yolov5s.pt --include onnx
    # step2: python -m onnxsim yolov5s.onnx yolov5s_sim.onnx
    net=cv2.dnn.readNet(r'E:\pythonProject\yolov5master\yolov5s_sim.onnx')
    blob=cv2.dnn.blobFromImage(img,1/255.,shape,swapRB=True,crop=False)
    net.setInput(blob)
    pred=net.forward()
    print(pred.shape)
    ids,confs,boxes=infer(img,pred,shape)
    et=time.time()
    print("run time:{:.2f}s/{:.2f}FPS".format(et-st,1/(et-st)))
    draw_rect(src,ids,confs,boxes)
 